#!/usr/bin/python
# -*- coding: utf-8 -*-
"""
@author coldwind
"""

import tensorflow as tf
import numpy as np


class TextDNN:
    def __init__(self,
                 model_type,
                 sequence_length,
                 number_class,
                 word_embedding,
                 vocab_size,
                 embedding_size,
                 hidden_layers,
                 hidden_size,
                 l2_reg_lambda=0.0):

        self.input_x = tf.placeholder(shape=[None, sequence_length], dtype=tf.int32, name='input_x')
        self.input_y = tf.placeholder(shape=[None, number_class], dtype=tf.float32, name='input_y')
        self.dropout_keep_prob = tf.placeholder(dtype=tf.float32, name='dropout_keep_prob')
        self.learning_rate = tf.placeholder(dtype=tf.float32, name='learning_rate')

        self.l2_loss = tf.constant(0.0, name='l2_loss')

        # embedding
        with tf.device('/cpu:0'), tf.name_scope('embedding'):
            if word_embedding:
                self.Embedding = tf.Variable(initial_value=word_embedding, dtype=tf.float32, name='word_embedding')
            else:
                self.Embedding = tf.get_variable(
                    initializer=tf.random_uniform(shape=[vocab_size, embedding_size], minval=-1, maxval=-1),
                    dtype=tf.float32, name='word_embedding')

            # (batch_size, sequence_length, embedding_size)
            self.embedded_words = tf.nn.embedding_lookup(self.Embedding, self.input_x, name='embedded_words')

            # 计算总特征数
            total_feature_counts = sequence_length * embedding_size
            # (batch_size, total_feature_counts)
            x = tf.reshape(self.embedded_words, shape=[-1, total_feature_counts])

        # Create fully-connected layers
        with tf.name_scope('full-connect'):
            def fc(x, num_hidden_units, name):
                with tf.variable_scope(name):
                    in_dim = x.shape[-1].value
                    value = 1.0 / np.sqrt(in_dim)
                    w = tf.get_variable(initializer=tf.random_uniform(shape=[in_dim, num_hidden_units], minval=-value,
                                                                      maxval=value), dtype=tf.float32, name='W')

                    b = tf.get_variable(initializer=tf.random_uniform(shape=[num_hidden_units], minval=-value,
                                                                      maxval=value), dtype=tf.float32, name='b')

                    output = tf.matmul(x, w) + b
                    return output

            for i in range(hidden_layers):
                x = tf.nn.elu(fc(x, hidden_size, name='l-%d' % i))

            self.output = x

        # dropout
        with tf.name_scope('dropout'):
            self.h_drop = tf.nn.dropout(self.output, keep_prob=self.dropout_keep_prob, name='dropout')

        # output
        with tf.name_scope('output'):
            dim = self.h_drop.shape[-1].value
            w = tf.get_variable(shape=[dim, number_class], dtype=tf.float32,
                                initializer=tf.contrib.layers.xavier_initializer(), name='W')
            b = tf.Variable(tf.constant(value=0.1, shape=[number_class]), name='b')

            self.l2_loss += tf.nn.l2_loss(w)
            self.l2_loss += tf.nn.l2_loss(b)

            self.scores = tf.nn.xw_plus_b(x=x, weights=w, biases=b, name='scores')
            if model_type == 'cls':
                self.predictions = tf.argmax(input=self.scores, axis=1, name='predictions')
            elif model_type == 'reg':
                self.predictions = tf.reduce_max(self.scores, axis=1, name='predictions')
                self.predictions = tf.expand_dims(self.predictions, -1)
        # loss
        with tf.name_scope('loss'):
            if 'cls' == model_type:
                # (batch_size,)
                losses = tf.nn.softmax_cross_entropy_with_logits_v2(labels=self.input_y, logits=self.scores)
                self.loss = tf.reduce_mean(losses) + self.l2_loss * l2_reg_lambda
            elif 'reg' == model_type:
                losses = tf.sqrt(tf.losses.mean_squared_error(labels=self.input_y, predictions=self.predictions))
                self.loss = tf.reduce_mean(losses) + self.l2_loss * l2_reg_lambda

        # accuracy
        with tf.name_scope('accuracy'):
            if 'cls' == model_type:
                correct_predictions = tf.equal(tf.argmax(self.input_y, axis=1), self.predictions)
                self.accuracy = tf.reduce_mean(tf.cast(correct_predictions, dtype=tf.float32), name='accuracy')
            elif 'reg' == model_type:
                pass


if __name__ == '__main__':
    text_dnn = TextDNN(model_type='cls',
                       sequence_length=10,
                       number_class=2,
                       word_embedding=None,
                       vocab_size=100,
                       embedding_size=128,
                       hidden_layers=5,
                       hidden_size=64,
                       l2_reg_lambda=0.1)
